import csv
import sys
import os
import requests
from aip import AipOcr
from sklearn.cluster import KMeans
import numpy as np
from PIL import ImageGrab
import re
import xlrd
import xlsxwriter

import cv2
# from matplotlib.pyplot import plot as plt
import matplotlib.pyplot as plt

# APP_ID = '18518188'  # '23555110'
# API_KEY = 'nDf9rK3oaKXdX6jIRuITMDOo'  # 'wZwS1tBPt1CBeqZwV8ECgOFz'
# SECRET_KEY = 'Tuj9iy9UADVuyHoLkCohAWwhvfT6z8Il'  # 'ZsTj1QzzFWck3GfENgGwRo4KGtOcIsGP'
APP_ID = '23555110'
API_KEY = 'wZwS1tBPt1CBeqZwV8ECgOFz'
SECRET_KEY = 'ZsTj1QzzFWck3GfENgGwRo4KGtOcIsGP'

options = {"language_type": "CHN_ENG", "detect_direction": "true", "detect_language": "true",
           "probability": "false",'table_border':'none','is_sync':'true','request_type':'excel'}


def get_file_content(filePath):
    with open(filePath, 'rb') as fp:
        return fp.read()


def save_as_csv(data, n_col, file_name='test.csv', mode='w'):
    """
    save data to a .csv file
    :param data: raw data in order, line by line, from left to right
    :param n_col: the number of columns
    :return: the number of rows
    """
    row = [None] * 8
    now = 0
    n = 0
    with open(file_name, mode)as f:#, encoding='utf-8'
        f_csv = csv.writer(f)
        for words in data:
            w = words['words']
            row[now] = w
            now += 1
            if now >= n_col:
                f_csv.writerow(row)
                now = 0
                n += 1
    return n


def save_OCR(url, dst='data_0.csv', mode='w'):
    r = requests.get(url)

    if mode=='a': # append mode
        temp = 'temp.xls'
        with open(temp, 'w' + 'b') as f:
            f.write(r.content)
        append_excel_to(temp, dst)
    elif mode == 'w':
        with open(dst, 'w'+'b') as f:
            f.write(r.content)


def to_OCR(path, dst='data_0.csv', mode='w'):
    client = AipOcr(APP_ID, API_KEY, SECRET_KEY)

    image = get_file_content(path)
    # result = client.basicGeneral(image, options)
    # result = client.accurate(image, options)
    result = client.tableRecognition(image,options)
    save_OCR(result['result']['result_data'],dst,mode)
    # result = client.getTableRecognitionResult(result['result']['request_id'])
    return None
    '''img = cv2.imread(path)
    # print(result['words_result'])
    xs = []
    ys = []
    ws = []
    hs = []
    i = 0
    for words in result['words_result']:
        loc = words['location']
        x = loc['left']
        y = loc['top']
        w = loc['width']
        h = loc['height']
        xs.append(x)
        ys.append(y)
        ws.append(w)
        hs.append(h)

        cv2.rectangle(img, (x, y), (x + w, y + h), (255, 255, 0))
        i += 1
        cv2.putText(img, str(i), (x, y), cv2.QT_FONT_BLACK, 0.5, (0, 0, 255))
    cv2.imshow('img', img)
    cv2.waitKey(1000)
    # plt.bar(np.arange(len(ys)),ys)
    xs = np.array(xs).reshape(-1, 1)
    ys = np.array(ys).reshape(-1, 1)

    # plt.show()
    km = KMeans()
    km.fit(xs)
    # print(km.cluster_centers_)
    n = km.n_clusters
    km.fit(ys)
    m = km.n_clusters

    rn = save_as_csv(result['words_result'], n, dst, mode)
    print('add%d rows' % rn)'''


def path_slove(fname='data_', ftype='.csv'):
    '''
    data_(\d+).csv
    :return:
    '''
    fs = os.listdir('.\\')
    fs = filter(lambda f: f.endswith(ftype), fs)
    max_n = 0
    last = None
    for f in fs:
        res = re.findall(fname + '(\d+)' + ftype, f)
        if len(res) < 1: continue
        n = int(res[0])
        if n >= max_n:
            last = f
            max_n = n + 1
    fname = fname.replace('*', 'default_name')
    return last, fname + str(max_n) + ftype


def main(arg):
    '''
    [-l] to save to last file, overwrite it
    [-a] append to file
    [file_name] to save as file_name
    :param arg:
    :return:
    '''
    length = len(arg) - 1
    last, dst = path_slove('data_', '.xls')
    lasts, source = path_slove(' *', '.png')
    mode = 'w'
    args = []
    source_input = False
    for a in arg:
        if '-l' in a:
            dst = last
            length -= 1
            continue
        if '-a' in a:
            mode = 'a'
            length -= 1
            continue
        if a.endswith('.xls'):
            dst = a
            continue
        if a.endswith('.png'):
            source = a
            source_input = True
        args.append(a)

    if not source_input:
        img1 = ImageGrab.grabclipboard()
        if img1 is None:
            print('no img detected,search path')
            source = lasts
            if lasts is None:
                print('no photo source found')
                exit(1)
        else:
            img1.save(source)

    print('using source:', source)
    print('mode:', mode)
    to_OCR(source, dst, mode)
    print('to file:', dst)
    print('done')


def append_excel_to(source,dst):
    fh = xlrd.open_workbook(dst)
    table = fh.sheets()[0]
    n = table.nrows
    datavalue = []
    for row in range(n):
        rdata = table.row_values(row)
        datavalue.append(rdata)

    fh = xlrd.open_workbook(source)
    table = fh.sheets()[0]
    n = table.nrows
    for row in range(n):
        rdata = table.row_values(row)
        datavalue.append(rdata)
    os.remove(source)

    wb1 = xlsxwriter.Workbook(dst)
    ws = wb1.add_worksheet()
    print('read merge')
    for a in range(len(datavalue)):
        for b in range(len(datavalue[a])):
            c = datavalue[a][b]
            ws.write(a, b, c)
    wb1.close()

if __name__ == '__main__':
    '''
    using cmd: python main.py -l -a data_x.csv
    or: 
    import os
    os.system('python main.py -l -a data_x.csv')
    '''
    main(sys.argv)
